import math
import random

def assign_cluster(x, c):
    """
    将样本x分配到最近的质心c
    
    参数:
    x: 一个数据点 (列表或元组)
    c: 质心列表 [c1, c2, ..., ck]，每个质心是一个与x维度相同的点
    
    返回:
    cluster_index: 最近质心的索引
    min_distance: 到最近质心的距离
    """
    min_distance = float('inf')
    cluster_index = -1
    
    for i, centroid in enumerate(c):
        # 计算欧几里得距离
        distance = 0.0
        for j in range(len(x)):
            distance += (x[j] - centroid[j]) ** 2
        distance = math.sqrt(distance)
        
        if distance < min_distance:
            min_distance = distance
            cluster_index = i
    
    return cluster_index, min_distance

def Kmeans(data, k, epsilon=1e-4, iteration=100):
    """
    K-means聚类算法实现
    
    参数:
    data: 数据集，列表的列表，每个内层列表代表一个数据点
    k: 聚类数量
    epsilon: 收敛阈值，当质心变化小于此值时停止迭代
    iteration: 最大迭代次数
    
    返回:
    centroids: 最终质心列表
    clusters: 每个数据点所属的簇索引
    distances: 每个数据点到其质心的距离
    """
    
    # 1. 初始化：随机选择k个点作为初始质心
    n = len(data)
    if n < k:
        raise ValueError("数据点数量不能小于聚类数量k")
    
    # 随机选择k个不重复的索引
    indices = random.sample(range(n), k)
    centroids = [data[i][:] for i in indices]  # 深拷贝初始质心
    
    print(f"初始质心: {centroids}")
    
    # 存储每个点所属的簇和距离
    clusters = [-1] * n
    distances = [0.0] * n
    
    prev_centroids = None
    iter_count = 0
    
    for iter_count in range(iteration):
        print(f"\n=== 第 {iter_count + 1} 次迭代 ===")
        
        # 2. 分配步骤：将每个点分配到最近的质心
        cluster_assignments = [[] for _ in range(k)]
        
        for i, point in enumerate(data):
            cluster_idx, dist = assign_cluster(point, centroids)
            clusters[i] = cluster_idx
            distances[i] = dist
            cluster_assignments[cluster_idx].append(point)
        
        # 3. 更新步骤：重新计算质心
        new_centroids = []
        for i in range(k):
            if len(cluster_assignments[i]) == 0:
                # 如果簇为空，保持原质心
                new_centroids.append(centroids[i][:])
                print(f"簇 {i} 为空，保持原质心")
                continue
            
            # 计算新质心（均值）
            dimension = len(data[0])
            new_centroid = [0.0] * dimension
            
            for point in cluster_assignments[i]:
                for j in range(dimension):
                    new_centroid[j] += point[j]
            
            for j in range(dimension):
                new_centroid[j] /= len(cluster_assignments[i])
            
            new_centroids.append(new_centroid)
            print(f"簇 {i} 有 {len(cluster_assignments[i])} 个点，新质心: {[round(x, 4) for x in new_centroid]}")
        
        # 4. 检查收敛条件：质心变化是否小于epsilon
        max_shift = 0.0
        for i in range(k):
            shift = 0.0
            for j in range(len(centroids[i])):
                shift += (centroids[i][j] - new_centroids[i][j]) ** 2
            shift = math.sqrt(shift)
            max_shift = max(max_shift, shift)
        
        print(f"最大质心移动: {max_shift:.6f}")
        
        if max_shift < epsilon:
            print(f"算法在 {iter_count + 1} 次迭代后收敛")
            break
        
        centroids = new_centroids
        prev_centroids = centroids[:]
    
    else:
        print(f"达到最大迭代次数 {iteration}，停止迭代")
    
    return centroids, clusters, distances

def calculate_wcss(data, centroids, clusters):
    """
    计算簇内平方和（Within-Cluster Sum of Squares）
    """
    wcss = 0.0
    for i, point in enumerate(data):
        centroid = centroids[clusters[i]]
        distance = 0.0
        for j in range(len(point)):
            distance += (point[j] - centroid[j]) ** 2
        wcss += distance
    return wcss

# 测试代码
if __name__ == "__main__":
    # 创建测试数据（二维点集，明显分为3个簇）
    test_data = [
        [1.0, 1.0], [1.1, 1.2], [1.2, 0.9], [0.9, 1.1],  # 簇0
        [4.0, 4.0], [4.1, 3.9], [3.9, 4.1], [4.2, 3.8],  # 簇1  
        [7.0, 1.0], [7.1, 1.2], [6.9, 0.8], [7.2, 1.1]   # 簇2
    ]
    
    print("测试数据:")
    for i, point in enumerate(test_data):
        print(f"点 {i}: {point}")
    
    # 运行K-means算法
    k = 3
    epsilon = 0.001
    max_iterations = 100
    
    print(f"\n开始K-means聚类，k={k}, epsilon={epsilon}")
    centroids, clusters, distances = Kmeans(test_data, k, epsilon, max_iterations)
    
    # 输出结果
    print(f"\n=== 最终结果 ===")
    print(f"质心:")
    for i, centroid in enumerate(centroids):
        print(f"  簇 {i}: {[round(x, 4) for x in centroid]}")
    
    print(f"\n数据点分配:")
    cluster_points = [[] for _ in range(k)]
    for i, point in enumerate(test_data):
        cluster_points[clusters[i]].append(point)
        print(f"点 {point} -> 簇 {clusters[i]}, 距离: {distances[i]:.4f}")
    
    print(f"\n各簇统计:")
    for i in range(k):
        print(f"簇 {i}: {len(cluster_points[i])} 个点")
    
    # 计算评估指标
    wcss = calculate_wcss(test_data, centroids, clusters)
    print(f"\n簇内平方和 (WCSS): {wcss:.4f}")
    
    # 可视化结果（简单的文本可视化）
    print(f"\n=== 聚类结果可视化 ===")
    for i in range(k):
        print(f"簇 {i} 的点:")
        for point in cluster_points[i]:
            print(f"  {point}")